package midend;

import java.util.List;
import java.util.LinkedList;
import java.util.Queue;
import java.util.HashMap;
import java.util.ArrayList;
import java.util.Set;
import java.util.HashSet;

import midend.IR_Value.IR_FuncBlock;
import midend.IR_Value.IR_BaseBlock;
import midend.Loop;

public class LoopInfo {

    // 每个基本块所属的最内层循环
    private HashMap<IR_BaseBlock, Loop> loopInfo;
    private Set<Loop> topLevelLoops;
    private Set<Loop> Loops;

    public void loopAnalysis(IR_FuncBlock funcBlock) {
        loopInfo = new HashMap<>();
        topLevelLoops = new HashSet<>();
        Loops = new HashSet<>();

        Dom_Analysis domAnalysis = new Dom_Analysis(funcBlock);
        List<IR_BaseBlock> postOrder = domAnalysis.getPostOrder(funcBlock);

        for (IR_BaseBlock block : postOrder) {
            LinkedList<IR_BaseBlock> workList = new LinkedList<>();
            for (IR_BaseBlock preBlock : block.getPreBlock()) {
                // 找到回边，将回边源节点即latch节点加入workList
                if (domAnalysis.dominates(block, preBlock)) {
                    workList.add(preBlock);
                }
            }
            // 如果存在回边，以block为header构建loop
            if (workList.size() > 0) {
                Loop loop = new Loop(block);
                mapBlockToLoop(loop, workList);
                Loops.add(loop);
            }
        }

        for (IR_BaseBlock block : postOrder) {
            updateOuterLoopsAndTopLevel(block);
        }

        funcBlock.setLoopInfo(loopInfo);
        funcBlock.setTopLevelLoops(topLevelLoops);
        funcBlock.setLoops(Loops);
    }

    // 将basicBlock映射到对应的最内循环，并维护Loop之间的父子关系
    private void mapBlockToLoop(Loop loop, List<IR_BaseBlock> workList) {

        while (!workList.isEmpty()) {
            IR_BaseBlock block = workList.remove(0);
            Loop blkLoop = loopInfo.get(block);
            if (blkLoop == null) {
                // 对应最内层循环就是loop的结点
                loopInfo.put(block, loop);
                if (block != loop.getHeader()) {
                    workList.addAll(block.getPreBlock());
                }
            } else {
                // 处于内层循环的节点
                Loop outMostLoop = blkLoop;
                while (outMostLoop.getParentLoop() != null) {
                    outMostLoop = outMostLoop.getParentLoop();
                }

                if (outMostLoop == loop) {
                    continue;
                }

                outMostLoop.setParentLoop(loop);
                loop.addChildLoop(outMostLoop);

                // 对与内部循环，只需要有一个节点被处理就够了，不在outMostLoop中的节点。
                IR_BaseBlock outMostLoopHeader = outMostLoop.getHeader();
                for (IR_BaseBlock toAdd : outMostLoopHeader.getPreBlock()) {
                    if (!outMostLoop.isInLoop(toAdd)) {
                        workList.add(toAdd);
                    }
                }
            }
        }
    }

    // 将block加入到自己所在循环体中，同时设置topLevelLoops
    private void updateOuterLoopsAndTopLevel(IR_BaseBlock block) {
        Loop innerLoop = loopInfo.get(block);
        if (innerLoop == null) {
            return;
        }

        if (innerLoop.getParentLoop() == null) {
            topLevelLoops.add(innerLoop);
        }
        Loop loop = innerLoop;
        while (loop != null) {
            if (block != loop.getHeader()) {
                // header在创建时已经加入
                loop.addLoopBody(block);
            }
            loop = loop.getParentLoop(); // 向上遍历，直到最外层循环
        }

    }

}
